from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from qwen3.faiss_retriever import FaissRetriever
from typing import List, Tuple

# 初始化FastAPI应用
app = FastAPI(title="向量检索API", version="1.0")

# 添加CORS中间件，允许前端跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境中应指定具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化检索器
retriever = FaissRetriever(
    embedding_model_path="Qwen/Qwen3-Embedding-0.6B",
    reranker_model_path="Qwen/Qwen3-Reranker-0.6B",
    device="cpu"  # 根据实际环境选择cpu或cuda
)

# 加载FAISS索引和文档（确保索引文件存在）
try:
    retriever.load_index("faiss_index.bin", "documents.txt")
except Exception as e:
    raise HTTPException(status_code=500, detail=f"索引加载失败: {str(e)}")

# 定义请求模型
class QueryRequest(BaseModel):
    query: str
    top_k_recall: int = 100
    top_k_rerank: int = 5

# 定义响应模型
class RetrievalResponse(BaseModel):
    results: List[Tuple[str, float]]

# 定义检索接口
@app.post("/retrieve", response_model=RetrievalResponse)
async def retrieve(request: QueryRequest):
    try:
        # 调用检索器获取结果
        results = retriever.retrieve_and_rerank(
            query=request.query,
            top_k_recall=request.top_k_recall,
            top_k_rerank=request.top_k_rerank
        )
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"检索失败: {str(e)}")

# 启动服务器
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)